from dotenv import load_dotenv
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime
import os
from sqlalchemy.orm import declarative_base, sessionmaker
from contextlib import contextmanager
from sqlalchemy.inspection import inspect


Base = declarative_base()

load_dotenv()  # 默认加载当前目录下的.env文件


class Video(Base):
    __tablename__ = "douyin_videos"
    id = Column(Integer, primary_key=True)
    video_path = Column(String(200), comment="视频本地路径")
    title = Column(String(100), comment="视频标题")
    create_time = Column(DateTime(), comment="下载时间")


class Audio(Base):
    __tablename__ = "audio_files"
    id = Column(Integer, primary_key=True)
    video_id = Column(Integer, comment="外键，关联视频")
    audio_path = Column(String(200), comment="音频本地路径")
    cos_url = Column(String(200), comment="腾讯云COS地址")
    create_time = Column(DateTime(), comment="创建时间")


class Content(Base):
    __tablename__ = "text_contents"
    id = Column(Integer, primary_key=True)
    audio_id = Column(Integer, comment="外键，关联音频")
    raw_text = Column(Text(), comment="原始识别文本")
    platform = Column(String(50), comment="平台（如'知乎'、'小红书'）")
    title = Column(String(100), comment="标题")
    intro = Column(Text(), comment="正文")
    tags = Column(String(200), comment="标签（逗号分隔）")
    create_time = Column(DateTime(), comment="创建时间")


class Image(Base):
    __tablename__ = "images"
    id = Column(Integer, primary_key=True)
    text_id = Column(Integer, comment="外键，关联文本内容")
    image_path = Column(String(300), comment="图片本地路径")
    create_time = Column(DateTime(), comment="创建时间")


class Record(Base):
    __tablename__ = "publish_records"
    id = Column(Integer, primary_key=True)
    text_id = Column(Integer, comment="外键，关联文本内容")
    platform = Column(String(50), comment="平台（如'知乎'、'小红书'）")
    status = Column(String(20), comment="发布状态（成功/失败）")
    publish_time = Column(DateTime(), comment="发布时间")
    extra_info = Column(Text(), comment="额外信息（如返回的链接、错误信息等）")


MYSQL_HOST = os.getenv("MYSQL_HOST")
MYSQL_USER = os.getenv("MYSQL_USER")
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD")
MYSQL_PORT = os.getenv("MYSQL_PORT")
MYSQL_DB = os.getenv("MYSQL_DB")
print(f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}/{MYSQL_DB}")
# 创建引擎
engine = create_engine(
    f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}/{MYSQL_DB}"
)
# 创建表
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)


def model_to_dict(obj):
    return {c.key: getattr(obj, c.key) for c in inspect(obj).mapper.column_attrs}


@contextmanager
def get_db_session():
    session = Session()
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        raise
    finally:
        session.close()
